import cv2
import os
import numpy as np
from sklearn import preprocessing


# 定义一个类，来处理与类标签编码相关的所有任务
# 训练器训练只认np.array类型的数据
class LabelEncoder(object):
    def __init__(self):
        self.le = None

    # 3.定义一个方法来为这些标签编码。在输入训练数据中，标签用单词表示，但我们需要数字来训练系统。
    # 该方法将定义一个预处理对象，该对象将单词转换成数字，同时保留这种前向后向的映射关系。
    # 将单词转换为数字的编码方法
    def encode_labels(self, label_words):
        self.le = preprocessing.LabelEncoder()
        self.le.fit(label_words)

    # 4.定义一个将单词转换成数字的方法
    # 将输入单词转换成数字
    def word_to_num(self, label_word):
        return int(self.le.transform([label_word])[0])

    # 5.定义一个方法，用于将数字转换回原始单词
    # 将数字转换成单词
    def num_to_word(self, label_num):
        return self.le.inverse_transform([label_num])[0]


def get_images_and_labels(path):
    label_words = []
    image_paths = []
    for root, dirs, files in os.walk(path):
        for filename in (x for x in files if x.endswith('.jpg')):
            filepath = os.path.join(root, filename)
            image_paths.append(filepath)
            label_words.append(filepath.split('/')[-2])
    return label_words, image_paths


def train_face_LBPH(pic_path, model_path):
    # 1.图片集目录和模型保存目录
    # path = "picture/train"
    # model_path = "dataset"

    # 2.定义模型文件路径和模型预测文件路径
    model_predicted_path = model_path + "/predicted.txt"
    model_file_path = model_path + "trainer.yam"

    # 3. 返回结果
    result = "人脸训练模型路径：{}\n".format(model_predicted_path)
    result += "人脸预测标签文件：{}\n".format(model_file_path)

    # 4.得到标签和图片路径
    label_words, image_paths = get_images_and_labels(pic_path)

    # 5.初始化变量
    le = LabelEncoder()
    le.encode_labels(label_words)
    images = []
    labels = []
    # DNN深度学习训练网络
    net = cv2.dnn.readNetFromCaffe("dataset/DNN_model/deploy.prototxt",
                                   "dataset/DNN_model/res10_300x300_ssd_iter_140000_fp16.caffemodel")
    # 6.循环训练每一张图片
    for i in range(len(label_words)):
        # print(label_words[i], image_paths[i])
        img = cv2.imread(image_paths[i])
        # 7.LBPH模型训练需要输入灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        (h, w) = img.shape[:2]

        result += "正在训练：{},{}\n".format(label_words[i], image_paths[i])
        # print("正在训练：{},{}".format(label_words[i], image_paths[i]))

        blob = cv2.dnn.blobFromImage(img, 1.0, (300, 300), [104, 177, 123], False, False)
        net.setInput(blob)
        detections = net.forward()
        # 8.人脸检测
        for j in range(0, detections.shape[2]):
            confidence = detections[0, 0, j, 2]
            if confidence > 0.7:
                box = detections[0, 0, j, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype('int')
                # face_w = endX - startX
                # face_h = endY - startY
                # cv2.rectangle(img, (startX, startY), (endX, endY), (255, 0, 0), 3)
                images.append(gray[startY:endY, startX:endX])
                labels.append(le.word_to_num(label_words[i]))
                # cv2.imshow('image', img)
                # cv2.waitKey(300)
    # print(len(images), len(labels))

    # 4.标签去重，用于输出预测文件
    label = sorted(set(label_words), key=label_words.index)
    label_num = sorted(set(labels), key=labels.index)
    with open(model_predicted_path, 'w') as f:
        for i in range(len(label)):
            f.write(str(label_num[i]) + ':' + str(label[i]) + '\n')
    result += "人脸标签预览：{}\n".format(label)

    # 9.创建LBPH人脸识别对象 训练、保存
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.train(images, np.array(labels))
    recognizer.save(model_file_path)

    result += "训练完成\n"
    # print("训练完成")
    return result


get_images_and_labels("picture/train")
pic_paths = 'picture/train'
model_paths = 'dataset'
train_face_LBPH(pic_paths, model_paths)
